import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from torch.utils.data import Dataset, DataLoader

# 设置随机种子以确保结果可重现
torch.manual_seed(42)
np.random.seed(42)


# 读取数据
def read_data(csv_path):
    # 读取CSV文件，处理可能的编码问题
    try:
        data = pd.read_csv(csv_path, encoding='utf-8-sig')
        print(f"成功读取数据：{csv_path}，共 {len(data)} 行记录")
        return data
    except Exception as e:
        print(f"读取数据错误 {csv_path}: {e}")
        return pd.DataFrame()  # 返回空DataFrame以防止后续错误


# 预处理函数
def preprocess_data(data, income_type):
    # 筛选特定类型的收入数据
    df = data[data['指标'] == income_type].copy()

    # 检查筛选后的数据是否为空
    if df.empty:
        print(f"警告：未找到指标为 '{income_type}' 的数据记录")
        return None, None, None

    # 过滤数据值为空的行
    df = df.dropna(subset=['数据值'])
    if df.empty:
        print(f"警告：指标 '{income_type}' 的有效数据值为空")
        return None, None, None

    # 使用映射将中文季度转换为数字
    quarter_mapping = {'第一季度': 1, '第二季度': 2, '第三季度': 3, '第四季度': 4}
    df['季度数字'] = df['季度'].map(quarter_mapping)

    # 过滤季度映射失败的数据
    df = df.dropna(subset=['季度数字'])
    if df.empty:
        print(f"警告：指标 '{income_type}' 没有有效的季度数据")
        return None, None, None

    # 创建时间字符串
    df['时间'] = df['年份'].astype(str) + 'Q' + df['季度数字'].astype(int).astype(str)

    # 确保按时间排序
    df = df.sort_values(by=['年份', '季度数字'])

    # 提取数值序列
    values = df['数据值'].values.reshape(-1, 1)

    # 数据标准化
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_values = scaler.fit_transform(values)

    return scaled_values, scaler, df[['时间', '数据值']]


# 创建时间序列数据集
def create_dataset(data, look_back=4):
    # 检查数据量是否足够
    if len(data) < look_back + 1:
        print(f"警告：数据量不足，需要至少 {look_back + 1} 个样本，实际只有 {len(data)} 个")
        return None, None

    X, y = [], []
    for i in range(len(data) - look_back):
        X.append(data[i:(i + look_back), 0])
        y.append(data[i + look_back, 0])
    return np.array(X), np.array(y)


# 定义PyTorch数据集
class TimeSeriesDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.float32)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


# 定义LSTM模型
class LSTMModel(nn.Module):
    def __init__(self, input_size=1, hidden_layer_size=50, output_size=1):
        super().__init__()
        self.hidden_layer_size = hidden_layer_size
        self.lstm = nn.LSTM(input_size, hidden_layer_size, batch_first=True)
        self.linear = nn.Linear(hidden_layer_size, output_size)
        self.hidden_cell = None

    def forward(self, input_seq):
        input_seq = input_seq.view(len(input_seq), -1, 1)
        lstm_out, self.hidden_cell = self.lstm(input_seq)
        predictions = self.linear(lstm_out[:, -1])
        return predictions


# 训练函数
def train_model(model, train_loader, criterion, optimizer, epochs=300):
    model.train()
    for epoch in range(epochs):
        total_loss = 0.0
        for seq, labels in train_loader:
            optimizer.zero_grad()
            y_pred = model(seq)
            loss = criterion(y_pred, labels.unsqueeze(1))
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        # 每50个epoch打印一次平均损失
        if epoch % 50 == 0:
            avg_loss = total_loss / len(train_loader)
            print(f'Epoch {epoch} 平均Loss: {avg_loss:.6f}')


# 预测函数
def predict_next_quarter(model, last_sequence, scaler):
    model.eval()
    with torch.no_grad():
        last_sequence = torch.tensor(last_sequence, dtype=torch.float32)
        prediction = model(last_sequence.unsqueeze(0))
        # 反标准化预测值
        prediction = scaler.inverse_transform(prediction.numpy().reshape(-1, 1))
    return prediction[0][0]


# 主程序
def main():
    # 定义要预测的收入类型
    income_types = {
        '居民': '居民人均可支配收入',
        '城镇': '城镇居民人均可支配收入',
        '农村': '农村居民人均可支配收入'
    }

    results = {}
    # 修复路径转义警告，使用原始字符串
    csv_path = r"F:\pycharm\python\pythonProject2\天津市居民人均可支配收入数据_20250712102415.csv"
    data = read_data(csv_path)

    # 打印数据中的所有指标，方便排查问题
    if not data.empty:
        print("\n数据中包含的指标：")
        for idx, indicator in enumerate(data['指标'].unique(), 1):
            print(f"{idx}. {indicator}")

    for name, income_type in income_types.items():
        print(f"\n===== 处理 {name}收入数据 =====")

        # 预处理数据
        scaled_values, scaler, history = preprocess_data(data, income_type)

        # 检查预处理结果
        if scaled_values is None:
            print(f"{name}收入数据处理失败，跳过该类型\n")
            continue

        # 创建训练数据集
        look_back = 4  # 使用过去4个季度预测下一个季度
        X, y = create_dataset(scaled_values, look_back)

        # 检查数据集是否有效
        if X is None or y is None:
            print(f"{name}收入数据不足，无法创建训练集，跳过该类型\n")
            continue

        # 创建数据加载器
        dataset = TimeSeriesDataset(X, y)
        # 动态调整batch_size，避免样本数不足
        batch_size = min(8, len(dataset))
        train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

        # 初始化模型
        model = LSTMModel()
        criterion = nn.MSELoss()
        optimizer = optim.Adam(model.parameters(), lr=0.001)

        # 训练模型
        print(f"开始训练{name}收入预测模型...")
        train_model(model, train_loader, criterion, optimizer, epochs=500)

        # 检查是否有足够的数据用于预测
        if len(scaled_values) < look_back:
            print(f"{name}收入数据不足，无法进行预测\n")
            continue

        # 获取最后4个季度的数据用于预测
        last_sequence = scaled_values[-look_back:].flatten()

        # 预测下一个季度
        prediction = predict_next_quarter(model, last_sequence, scaler)

        # 保存结果
        results[name] = {
            'prediction': prediction,
            'history': history,
            'model': model,
            'scaler': scaler
        }

        print(f"预测2024年第四季度{name}收入: {prediction:.2f}元\n")

    # 可视化结果（如果有有效结果）
    if results:
        plt.figure(figsize=(14, 8))
        plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体
        plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

        # 绘制历史数据和预测
        for name, result in results.items():
            history = result['history']
            prediction = result['prediction']

            # 历史数据
            plt.plot(history['时间'], history['数据值'], 'o-', label=f'{name}历史数据')

            # 预测点
            plt.plot('2024Q4', prediction, 's', markersize=10, label=f'{name}预测值')

        # 设置图表属性
        plt.title('天津市各类收入历史数据及2024年第四季度预测', fontsize=16)
        plt.xlabel('时间', fontsize=14)
        plt.ylabel('收入（元）', fontsize=14)
        plt.xticks(rotation=45)
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.legend()
        plt.tight_layout()  # 自动调整布局，防止标签被截断
        plt.show()

        # 打印预测结果表格
        print("\n===== 预测结果汇总 =====")
        print(f"{'收入类型':<12}{'2024年第四季度预测值(元)':<25}")
        print("-" * 40)
        for name, result in results.items():
            print(f"{name + '收入':<12}{result['prediction']:<25.2f}")
        print("-" * 40)
    else:
        print("\n没有有效的预测结果可展示")


if __name__ == "__main__":
    main()